import os
import cv2
import cv2 as cv
import json
import numpy as np
from multiprocessing.pool import  Pool
from collections import defaultdict

target_path = '/my_data'

surfGPU = cv.cuda.SURF_CUDA_create(10000, _nOctaves=4, _nOctaveLayers=3, _upright=True)
matcherGPU = cv.cuda.DescriptorMatcher_createBFMatcher(cv.NORM_L2)
std_height = 2048

def aling(img_name):
    print(img_name)
    im_path = os.path.join(img_path, img_name)
    im_rgb = cv2.imread(im_path)
    im = cv2.cvtColor(im_rgb, cv2.COLOR_BGR2GRAY)

    im_path = os.path.join(ori_temp_path, img_name.replace('.jpg', '_t.jpg'))
    rgb_temp = cv2.imread(im_path)
    im_t = cv2.cvtColor(rgb_temp, cv2.COLOR_BGR2GRAY)

    im_t = (im_t + (np.mean(im) - np.mean(im_t))).astype(np.uint8)

    height, width = im_t.shape[:2]
    h_ratio = 1.0 * height / std_height
    std_width = int(round(width / h_ratio))
    w_ratio = 1.0 * width / std_width


    im_ = cv2.resize(im, (std_width, std_height))
    im_ = cv2.equalizeHist(im_)
    im_t_ = cv2.resize(im_t, (std_width, std_height))
    im_t_ = cv2.equalizeHist(im_t_)


    gpu_im_ = cv.cuda_GpuMat(im_)
    gpu_im_t_ = cv.cuda_GpuMat(im_t_)
    kp1gpu, d1gpu = surfGPU.detectWithDescriptors(gpu_im_, None)
    kp2gpu, d2gpu = surfGPU.detectWithDescriptors(gpu_im_t_, None)
    kp1gpu = surfGPU.downloadKeypoints(kp1gpu)
    kp2gpu = surfGPU.downloadKeypoints(kp2gpu)


    matches_gpu = matcherGPU.match(d1gpu, d2gpu)
    matches_gpu = [m for m in matches_gpu if m.distance <= 0.7]
    no_of_matches = len(matches_gpu)

    if no_of_matches >= 4:
        # Define empty matrices of shape no_of_matches * 2.
        p1 = np.zeros((no_of_matches, 2))
        p2 = np.zeros((no_of_matches, 2))

        for i in range(len(matches_gpu)):
            p1[i, :] = kp1gpu[matches_gpu[i].queryIdx].pt
            p2[i, :] = kp2gpu[matches_gpu[i].trainIdx].pt

        p1 = p1 * [[w_ratio, h_ratio]]
        p2 = p2 * [[w_ratio, h_ratio]]


        homography, mask = cv2.findHomography(p2, p1, cv2.RANSAC)
        transformed_img = cv2.warpPerspective(rgb_temp, homography, (width, height))

    else:

        transformed_img = rgb_temp
    cv2.imwrite(os.path.join(fast_align_temp_path, img_name.replace('.jpg', '_t.jpg')), transformed_img)

    # res_1 = (im - np.mean(1.0 * im)) - (transformed_img - np.mean(1.0 * transformed_img))
    # res_1 = res_1 - np.min(res_1)
    # res_1 = (res_1 / np.max(res_1) * 255).astype(np.uint8)
    #
    # bbs = im_name_dic[img_name]
    # t= False
    # for bb in bbs:
    #     im = cv2.rectangle(im, (int(bb[0]-20), int(bb[1]-20)), (int(bb[2]+10), int(bb[3]+10)),
    #                   thickness=1, color=(0, 0, 255))
    #     res_1 = cv2.rectangle(res_1, (int(bb[0]-20), int(bb[1]-20)), (int(bb[2]+10), int(bb[3]+10)),
    #                   thickness=1, color=(0, 0, 255))
    #     t = True
    # if t:
    # cv2.imwrite(os.path.join('/tcdata/111', img_name), np.concatenate((im, res_1), axis=1))
    # cv2.imwrite(os.path.join('/tcdata/111', img_name.replace('.jpg', '_t.jpg')), transformed_img)

if __name__ == '__main__':
    img_path = '/my_data/ori_img'
    ori_temp_path = '/my_data/ori_temp'

    # img_path = '/tcdata/ori_img'
    # ori_temp_path = '/tcdata/ori_temp'
    #
    # anno = '/tcdata/train_annos.json'
    # anno = json.load(open( anno,'r'))
    # im_name_dic = defaultdict(list)
    #
    # for ann in anno:
    #
    #     name = ann['name']
    #     im_name_dic[name].append(ann['bbox'])


    if not os.path.exists(os.path.join(target_path,'fast_align_temp' )):
        os.mkdir(os.path.join(target_path,'fast_align_temp' ))
    fast_align_temp_path = os.path.join(target_path,'fast_align_temp' )

    pool = Pool(processes=20)

    files_img = os.listdir(img_path)

    pool.map(aling,files_img)


